import xml.etree.ElementTree as ET
import json
import requests

# 定义命名空间
NAMESPACES = {
    'cwe': 'http://cwe.mitre.org/cwe-7',
    'xhtml': 'http://www.w3.org/1999/xhtml',
    'xsi': 'http://www.w3.org/2001/XMLSchema-instance'
}

def translate_text(text, source_lang="en", target_lang="zh"):
    """调用翻译API将文本翻译为中文(模拟实现)"""
    try:
        # 实际使用时替换为真实翻译API调用
        url = "https://api.mymemory.translated.net/get"
        params = {"q": text, "langpair": f"{source_lang}|{target_lang}"}
        response = requests.get(url, params=params)
        return response.json()["responseData"]["translatedText"]
    except:
        return f"[译] {text}"

def convert_xhtml_to_text(xhtml_element):
    """
    将xhtml元素转换为纯文本格式的代码
    处理标签: <xhtml:div>、<xhtml:br>、<xhtml:i>等
    """
    def process_element(elem, indent_level=0):
        result = ""
        indent = "    " * indent_level
        
        # 1. 处理当前元素的文本内容
        if elem.text and elem.text.strip():
            result += indent + elem.text.strip()
            # 确保内容有换行
            if result and not result.endswith('\n'):
                result += "\n"
        
        # 2. 处理所有子元素
        for child in elem:
            # 根据子元素类型分别处理
            if child.tag.endswith('div'):
                # 检测样式缩进
                style = child.attrib.get("style", "")
                new_level = indent_level + 1 if "margin-left" in style else indent_level
                child_content = process_element(child, new_level)
                
                # 确保内容有换行
                if result and not result.endswith('\n'):
                    result += "\n"
                result += child_content
                
            elif child.tag.endswith('br'):
                # 换行并保持缩进
                result += "\n"
                if child.tail and child.tail.strip():
                    result += indent + child.tail
                # if result and not result.endswith('\n'):
                #     result += "\n"
                
            elif child.tag.endswith('i'):
                # 获取斜体元素中的所有文本（包括深层嵌套）
                def extract_i_text(i_elem):
                    text = i_elem.text or ""
                    for sub in i_elem:
                        if sub.text:
                            text += sub.text
                        if sub.tail:
                            text += sub.tail
                            # pass
                    return text.strip()
                
                # 添加为注释
                i_text = extract_i_text(child)
                if i_text:
                    result += indent + f"{i_text}"
            
            else:
                # 其他元素递归处理
                child_content = process_element(child, indent_level)
                if child_content:
                    if result and not result.endswith('\n'):
                        result += "\n"
                    result += child_content
        
        # 3. 处理尾随文本
        if elem.tail and elem.tail.strip():
            # 处理特殊字符转义
            tail_text = elem.tail.strip()
            tail_text = tail_text.replace("&gt;", ">").replace("&lt;", "<")
            
            # 如果是换行后的文本，确保缩进
            if result.endswith('\n'):
                result += indent + tail_text
            else:
                result += " " + tail_text
        
        return result
    
    result = process_element(xhtml_element).strip()
    print(result)
    print("---------------------------------------------------------------------------\n")
    return result



def parse_cwe_xml(xml_file, output_json):
    # 解析XML文件
    tree = ET.parse(xml_file)
    root = tree.getroot()
    
    # 创建结果列表
    weaknesses = []
    
    # 遍历所有Weakness元素
    for weakness_elem in root.findall(".//cwe:Weakness", NAMESPACES):
        weakness_data = {
            "vulnerability_name": "",
            "缺陷名称": "",
            "description": "",
            "缺陷描述": "",
            "代码示例": []
        }
        
        # 1. 提取漏洞名称
        name_attr = weakness_elem.get("Name")
        if name_attr:
            weakness_data["vulnerability_name"] = name_attr
            weakness_data["缺陷名称"] = translate_text(name_attr)
        
        # 2. 提取描述
        desc_elem = weakness_elem.find("cwe:Description", NAMESPACES)
        if desc_elem is not None and desc_elem.text:
            description = desc_elem.text.strip()
            weakness_data["description"] = description
            weakness_data["缺陷描述"] = translate_text(description)
        
        # 3. 处理示例代码
        examples_elem = weakness_elem.find("cwe:Demonstrative_Examples", NAMESPACES)
        if examples_elem is not None:
            for example_elem in examples_elem.findall("cwe:Demonstrative_Example", NAMESPACES):
                example_data = {
                    "example_description": "",
                    "示例描述": "",
                    "代码": ""
                }
                
                # 3.1 示例描述
                intro_elem = example_elem.find("cwe:Intro_Text", NAMESPACES)
                if intro_elem is not None and intro_elem.text:
                    intro_text = intro_elem.text.strip()
                    example_data["example_description"] = intro_text
                    example_data["示例描述"] = translate_text(intro_text)
                
                # 3.2 代码提取
                code_elem = example_elem.find("cwe:Example_Code", NAMESPACES)
                if code_elem is not None:
                    # 获取代码语言
                    lang = code_elem.get("Language", "")
                    
                    # 查找xhtml:div元素
                    xhtml_div = code_elem.find("xhtml:div", NAMESPACES)
                    
                    if xhtml_div is not None:
                        code_text = convert_xhtml_to_text(xhtml_div)
                        example_data["代码"] = f"// Language: {lang}\n{code_text}"
                
                weakness_data["代码示例"].append(example_data)
        
        weaknesses.append(weakness_data)
    
    # 保存为JSON文件
    with open(output_json, "w", encoding="utf-8") as f:
        json.dump(weaknesses, f, ensure_ascii=False, indent=2)
    
    return weaknesses

# 使用示例
if __name__ == "__main__":
    input_xml = r"language_C\658.xml"  # 替换为你的XML文件路径
    output_json = r"language_C\cwe_weaknesses_C.json"  # 输出JSON文件
    
    results = parse_cwe_xml(input_xml, output_json)
    print(f"成功处理 {len(results)} 个漏洞，结果已保存至 {output_json}")